from typing import Callable, Sequence, List
from autogen_agentchat.agents import BaseChatAgent
from autogen_agentchat.base import Response
from autogen_agentchat.conditions import MaxMessageTermination
from autogen_agentchat.messages import ChatMessage, TextMessage
from autogen_agentchat.teams import SelectorGroupChat
from autogen_agentchat.ui import Console
from autogen_core import CancellationToken
from autogen_ext.models.openai import OpenAIChatCompletionClient


# Update the message history.
model_client = OpenAIChatCompletionClient(model="qwen2.5:7b",
                                          model_info={
                                              "vision": False,
                                              "function_calling": True,
                                              "family": "Qwen3",
                                              "structured_output": True,
                                              "json_output": True,
                                          },
                                          api_key="ollama",
                                          base_url="http://127.0.0.1:11434/v1")

# In this example, we create an agent class that can perform simple arithmetic operations on a given integer. Then, we will use different instances of this agent class in SelectorGroupChat to convert a given integer to another integer by applying a series of arithmetic operations.
# The ArithmeticAgent class accepts an operator_func that takes an integer and returns an integer after applying the arithmetic operation. In its on_messages method, it applies the operator_func to the integer in the input message and returns a response with the result.

class ArithmeticAgent(BaseChatAgent):
    def __init__(self, name: str, description: str, operator_func: Callable[[int], int]) -> None:
        super().__init__(name, description=description)
        self._operator_func = operator_func
        self._message_history: List[ChatMessage] = []

    @property
    def produced_message_types(self) -> Sequence[type[ChatMessage]]:
        return (TextMessage,)

    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
        # Update the message history.
        # 注意：messages 列表有可能为空，这意味着该代理之前已被选中。
        self._message_history.extend(messages)
        # 解析最后一条消息中的数字。
        assert isinstance(self._message_history[-1], TextMessage)
        number = int(self._message_history[-1].content)
        # 对数字应用运算符函数。
        result = self._operator_func(number)
        # 创建包含结果的新消息。
        response_message = TextMessage(content=str(result), source=self.name)
        # Update the message history.
        self._message_history.append(response_message)
        # 返回响应。
        return Response(chat_message=response_message)

    async def on_reset(self, cancellation_token: CancellationToken) -> None:
        pass


async def run_number_agents() -> None:
    # 创建用于数字运算的代理。
    add_agent = ArithmeticAgent("add_agent", "给数字加 1。", lambda x: x + 1)
    multiply_agent = ArithmeticAgent("multiply_agent", "将数字乘以 2。", lambda x: x * 2)
    subtract_agent = ArithmeticAgent("subtract_agent", "从数字中减去 1。", lambda x: x - 1)
    divide_agent = ArithmeticAgent("divide_agent", "将数字除以 2 并向下取整。", lambda x: x // 2)
    identity_agent = ArithmeticAgent("identity_agent", "原封不动返回数字。", lambda x: x)
    # 终止条件是在 10 条消息后停止。
    termination_condition = MaxMessageTermination(10)
    # 创建一个选择器组聊天。
    selector_group_chat = SelectorGroupChat(
        [add_agent, multiply_agent, subtract_agent, divide_agent, identity_agent],
        model_client=model_client,
        termination_condition=termination_condition,
        allow_repeated_speaker=True,  # 允许同一代理多次发言，此任务必需。
        selector_prompt=(
            "可用角色:\n{roles}\n他们的工作描述:\n{participants}\n"
            "当前对话历史:\n{history}\n"
            "请为下一条消息选择最合适的角色，仅返回角色名称。"
        ),
    )

    # 使用给定任务运行选择器组聊天并流式传输响应。
    task: List[ChatMessage] = [
        TextMessage(content="应用这些运算将给定数字转换为 25。", source="user"),
        TextMessage(content="10", source="user"),
    ]
    stream = selector_group_chat.run_stream(task=task)
    await Console(stream)



# 使用asyncio.run()运行异步函数。
if __name__ == "__main__":
    import asyncio
    asyncio.run(run_number_agents())